# Main class file for the Prinz summer project at Ottawa University during the summer of 08.
from numpy import *
from scipy import *
from pylab import *

from genutils.io import *
from genutils.handy_functions import *
from Prinz_functions import *

import os
import cPickle

class Ccell:
    '''The main class for the cell model, it holds information about the cell and can
       be exercised to produce a voltage trace. '''
    def __init__(self,config_filename='cell_config'):
        # ----------------------------------------------------------------------------
        # --------------------- READ IN FROM CONFIGURATION FILES ---------------------
        # ----------------------------------------------------------------------------
        # read in the configuration file for this cell
        self.config_filename = config_filename
        self.temperature = from_config('temperature',config_filename)
        self.V_init = from_config('V_init',config_filename)
        self.channels = from_config('channels',config_filename)
        self.conductances = from_config('conductances',config_filename)
        self.area = from_config('area',config_filename)
        self.capacitance = from_config('capacitance',config_filename)
        self.tau_Ca = from_config('tau_Ca',config_filename)
        self.f = from_config('f',config_filename)
        self.Ca_0 = from_config('Ca_0',config_filename)
        self.Ca_valence = from_config('Ca_valence',config_filename)
        self.synaptic_type = from_config_string('synaptic_type',config_filename)
        self.synaptic_conductance = from_config('synaptic_conductance',config_filename)
        self.synaptic_onset = from_config('synaptic_onset',config_filename)
        self.synaptic_duration = from_config('synaptic_duration',config_filename)
        self.syn_E = from_config('synaptic_E',config_filename)

        # leak parameters
        self.leak_conductance = from_config('leak_conductance',config_filename)
        self.leak_E = from_config('leak_E',config_filename)

        # stimulus parameters
        self.stimulus_type = from_config_string('stimulus_type',config_filename)
        if 'step' in self.stimulus_type:
            self.stimulus_onset = from_config('stimulus_onset',config_filename)
            self.stimulus_duration = from_config('stimulus_duration',config_filename)
            self.stimulus_amplitude = from_config('stimulus_amplitude',config_filename)

        # read in the configuration file for the different ion channels
        self.channel_config_filename = from_config_string('channel_config_filename',config_filename)
        self.channels_dict = read_dict_float(self.channel_config_filename)
        self.ext_Ca = from_config('Ext_Ca',self.channel_config_filename)

        # make a dictionary of the channels with ID as index and contents of [prefix,type]
        self.available_channels = {}
        for key in self.channels_dict.keys():
            if 'ID' in key:
                self.available_channels[self.channels_dict[key]] = [key[:-3]]
        # add the channel types to the available channels array
        for ID in self.available_channels.keys():
            channel_prefix = self.available_channels[ID][0]
            self.available_channels[ID].append(self.channels_dict[channel_prefix + '_type'])

        # ----------------------------------------------------------------------------
        # --------------------- SET UP ION CHANNELS ----------------------------------
        # ----------------------------------------------------------------------------
        self.set_channels = {}
        for ID in self.channels:
            # determine the type of channel it is
            if ID in self.available_channels.keys():
                self.set_channel(ID) # where we add to self.set_channels
            else:
                print self.available_channels
                raise RuntimeError("The channel id number you listed in the cell config file ("\
                                     + ID +") does not exist, you idiot!.")

        # check for the existance of any Calcium channels, so we can ensure that if it has any it has
        #   the infrastructure to determine the calcium concentration in the cell.
        self.any_Ca_channels = False
        for ID in self.set_channels.keys():
            if 'Ca' in self.ID_has_prefix(ID):
                self.any_Ca_channels = True

        # check to see if CaT and CaS are set.
        if self.any_Ca_channels:
            CaT_set = False
            CaS_set = False
            for ID in self.set_channels.keys():
                if 'CaT' in self.ID_has_prefix(ID):
                    CaT_set = True
                elif 'CaS' in self.ID_has_prefix(ID):
                    CaS_set = True
        
        # raise an error if the channels aren't set and need to be.
        if self.any_Ca_channels and not (CaT_set or CaS_set):
            raise RuntimeError('Calcium channels CaT and CaS must be set and are not both set.\nAlter the cell_config or channel_config file.')

        # ----------------------------------------------------------------------------
        # --------------------- SET UP FUNCTION FOR THE INTEGRATOR -------------------
        # ----------------------------------------------------------------------------
        # set up a dictionary of lists to store the currents
        self.currents = {}
        self.V_funct_terms = []
        # set up the function to be passed to the integrator
        X_size = 2 
        # if we don't have any calcium channels, don't bother calculating the calcium concentration.
        if not self.any_Ca_channels:
            X_size -= 1
        # X_size is the size of the X array that keeps track of all the dynamical variables
        #   X[0] = V
        #   X[1] = internal Ca concentration (If needed)
        #   X[2+] = m's and h's of the various currents   
        for ID in self.set_channels.keys():
            # what type of channel is this
            type = self.ID_is_type(ID)
            #---------- type 1 or type 2 ---------------------
            if type == 1 or type == 2:
                # keep track of what indexes of X correspond to the dynamical variables in this channel
                self.set_channels[ID]['m_index'] = X_size
                self.set_channels[ID]['h_index'] = X_size + 1
                # add 2 to X_size counter
                X_size += 2
            #---------- type 3 or type 4 (no h variable) ---------------------
            if type == 3 or type == 4:
                E_rev = self.channels_dict[prefix+'_E']
                # keep track of what indexes of X correspond to the dynamical variables in this channel
                self.set_channels[ID]['m_index'] = X_size
                # add 1 to X_size counter
                X_size += 1
            #---------- type 5 (E_rev is determined by [Ca+]) ---------------------
            if type == 5:
                # keep track of what indexes of X correspond to the dynamical variables in this channel
                self.set_channels[ID]['m_index'] = X_size
                self.set_channels[ID]['h_index'] = X_size + 1
                # add 2 to X_size counter
                X_size += 2
            # keep track of where the calcium currents lambda functions are.
            prefix = self.ID_has_prefix(ID)
            self.currents[prefix] = {'funct':self.set_currents(ID),'trace':array([]),'evaluations':[]}
        self.X_size = X_size
        self.total_time = 0
        self.t = array([])

    def set_currents(self,ID):
        prefix = self.ID_has_prefix(ID)
        type = self.ID_is_type(ID)
        g = self.conductances[self.channels.index(ID)]*1000 # to get the result into uS so the units work out
        A = self.area
        p = self.channels_dict[prefix+'_p']
        if type == 1 or type == 2:
            m_index = self.set_channels[ID]['m_index']
            h_index = self.set_channels[ID]['h_index']
            E_rev = self.channels_dict[prefix+'_E']
            return lambda X: A * g * X[m_index]**p * X[h_index] * (X[0] - E_rev)
        if type == 3 or type == 4:
            m_index = self.set_channels[ID]['m_index']
            E_rev = self.channels_dict[prefix+'_E']
            return lambda X: A * g * X[m_index]**p * (X[0] - E_rev)
        if type == 5:
            m_index = self.set_channels[ID]['m_index']
            h_index = self.set_channels[ID]['h_index']
            return lambda X: A * g * X[m_index]**p * X[h_index] * (X[0] - Nernst(self.temperature,self.Ca_valence,self.ext_Ca,X[1]))
            
    def V_trace(self):
        # just returns the voltage trace from the results of the exercising.
        return self.X.T[0]

    def Ca_trace(self):
        # just returns the trace of the intracellular Ca concentration from the results of the exercising.
        if self.any_Ca_channels:
            return self.X.T[1]
        else:
            raise RuntimeError('No Ca concentrations were calculated, because no Ca channels were specified.')

    def g_factor_trace(self,prefix):
        # returns a trace of the factor multiplying the maximum conductance of a current from the results of the exercising.
        for ID in self.set_channels.keys():
            prefix = self.ID_has_prefix(ID)
            p = self.channels_dict[prefix+'_p']
            m_index = self.set_channels[ID]['m_index']
            if 'h_index' in self.set_channels[ID].keys():
                h_index = self.set_channels[ID]['h_index']
                return self.X.T[m_index]**p * self.X.T[h_index]
            else:
                return self.X.T[m_index]**p
                

    def dX_dt(self,X,t):
        # the main function which returns the RHS (right hand side) of all the differential equations
        return_list = []
        # ----------- V -------------
        dV_dt = self.I_injected(t)

        for current_name in self.currents.keys():
            current = self.currents[current_name]['funct'](X)
            if 'CaT' in current_name:
                I_CaT = current
            elif 'CaS' in current_name:
                I_CaS = current
            dV_dt -= current
            self.currents[current_name]['evaluations'].append([t,current])
            
        dV_dt -= self.area * self.leak_conductance * 1000  * (X[0] - self.leak_E) # 1000 is to get conductance in units of uS 
        dV_dt -= self.area * (self.synaptic_input(t)) * (X[0] - self.syn_E)  # 1000 is to get conductance in units of uS
        dV_dt /= self.capacitance
        return_list.append( dV_dt )

        #----------- internal Ca concentration --------
        if self.any_Ca_channels:
            dCa_dt = (-self.f*(I_CaT + I_CaS) - X[1] + self.Ca_0)/self.tau_Ca
            return_list.append( dCa_dt )

        #----------- the rest of the dynamical variables ---------------- 
        for ID in self.set_channels.keys():
            type = self.ID_is_type(ID)
            V = X[0]
            # ---- add both an m and an h variable
            if type == 1 or type == 2 or type == 5:
                # add m
                m_funct = self.set_channels[ID]['m_funct']
                m_in_X_index = self.set_channels[ID]['m_index']
                m = X[m_in_X_index]
                return_list.append( m_funct(m,V) )
                # add h
                h_funct = self.set_channels[ID]['h_funct']
                h_in_X_index = self.set_channels[ID]['h_index']
                h = X[h_in_X_index]
                return_list.append( h_funct(h,V) )
            # ----- add an m with m handeled differently
            elif type == 3:
                # add m
                m_funct = self.set_channels[ID]['m_funct']
                m_in_X_index = self.set_channels[ID]['m_index']
                m = X[m_in_X_index]
                return_list.append( m_funct(m,V) )
                # no h
            else:
                # add m
                Ca = X[1]
                m_funct = self.set_channels[ID]['m_funct']
                m_in_X_index = self.set_channels[ID]['m_index']
                m = X[m_in_X_index]
                return_list.append( m_funct(Ca,m,V) )
                # no h
        return return_list
        
    def ID_has_prefix(self,ID):
        # make a maping of ID numbers to name prefixes
        return self.available_channels[ID][0]

    def ID_is_type(self,ID):
        # make a maping of ID numbers to types
        return self.available_channels[ID][1]

    def archive(self,directory):
        ''' This method just cPickles the cell into a directory given '''
        # see if the directory already exists
        listing = os.listdir(".")
        dir_exists = False
        for file in listing:
            if file == directory:
                dir_exists = True
        # if it doesn't exist, make it
        if not dir_exists:
            os.system("mkdir -p "+ directory)
        # archive the cell.
        # make a list of the important things you need to keep around.
        archive_dict = {}
        archive_dict['X'] =self.X
        archive_dict['t']=self.t
        archive_dict['dt']=self.dt
        archive_dict['total_time']=self.total_time
        archive_dict['config_filename']=self.config_filename
        archive_dict['currents'] = {}
        for key in self.currents.keys():
            archive_dict['currents'][key] = self.currents[key]['trace']
        cPickle.dump(archive_dict,open(directory+'/cell_data.cPickle','w'),True)

    # a way to read in a data archive of the cell.
    @classmethod
    def from_data_file(self,data_filename):
        data = cPickle.load(open(data_filename))
        c = Ccell(config_filename=data['config_filename'])
        c.X = data['X']
        c.t = data['t']
        c.dt = data['dt']
        c.total_time = data['total_time']
        for key in c.currents.keys():
            c.currents[key]['trace'] = data['currents'][key]
        c.config_filename = data['config_filename']
        return c

        
    def exercise(self,X_init=array([0,0]),timestep=0.1,i_time=1000,Prinz_initialize=True):
        ''' This method just causes the cell to be simulated for a time equal to i_time (ms)
            and with timesteps of size timestep. '''
        # set up the time vector
        X_was_passed = X_init.all()
        timesteps = i_time/timestep
        self.dt = timestep
        t = linspace(0,i_time,timesteps)
        # set the initialization values
        if not X_was_passed:
            X_init = self.initialize_X(Prinz_initialize=Prinz_initialize)

        # clear any previous evaluations of the current functions
        for prefix in self.currents.keys():
            self.currents[prefix]['evaluations'] = []

        # run the integrator
        if X_was_passed:
            self.X = array_extend(self.X,integrate.odeint(self.dX_dt, X_init, t))
            self.t = array_extend(self.t,t+self.total_time)
            self.total_time += i_time
        else:
            self.X = integrate.odeint(self.dX_dt,X_init,t)
            self.t = t
            self.total_time = i_time
        
        for prefix in self.currents.keys():
            # resample the evaluated currents to the grid provided by t
            evals = array(self.currents[prefix]['evaluations'])
            last_total_time = self.total_time - i_time
            c_vals = evals.T[1]
            t_vals = evals.T[0] + self.total_time - i_time 
            from scipy import interpolate
            new_c_vals_funct = interp1d(t_vals,c_vals)
            if X_was_passed:
                self.currents[prefix]['trace'] = array_extend(self.currents[prefix]['trace'],new_c_vals_funct(t+last_total_time)) 
            else:
                self.currents[prefix]['trace'] = new_c_vals_funct(t) 

    def initialize_X(self,Prinz_initialize=True):
        V_init = self.V_init
        Ca_0 = self.Ca_0
        X = range(self.X_size) # just setting the size of X
        X[0] = V_init
        if self.any_Ca_channels:
            X[1] = Ca_0
        for ID in self.set_channels.keys():
            type = self.ID_is_type(ID)
            # set the initial value for the m variable to m_inf(V_init)
            m_index = self.set_channels[ID]['m_index']
            m_inf = self.set_channels[ID]['m_inf']
            if Prinz_initialize:
                X[m_index] = 0
            elif type == 4: 
                X[m_index] = m_inf(Ca_0,V_init)
            else:
                X[m_index] = m_inf(V_init)

            # set the initial value for the h variable to h_inf(V_init)
            if 'h_index' in self.set_channels[ID].keys():
                h_index = self.set_channels[ID]['h_index']
                h_inf = self.set_channels[ID]['h_inf']
                if Prinz_initialize:
                    X[h_index] = 1
                else:
                    X[h_index] = h_inf(V_init)
        return X
    
    def I_injected(self,t):
        # this function describes the time dependence of the imput current
        if 'step' in self.stimulus_type:
             if t < self.stimulus_onset:
                return 0
             elif t < (self.stimulus_onset + self.stimulus_duration):
                return self.stimulus_amplitude
             else:
                return 0

    def synaptic_input(self,t):
        #this function describes the time dependence of the synaptic input
        if 'inhibit' in self.synaptic_type:
            if t < self.synaptic_onset:
                return 0
            elif t < (self.synaptic_onset + self.synaptic_duration):
                return self.synaptic_conductance
            else:
                return 0

    def set_synaptic_time(self,onset,duration):
        self.synaptic_onset = onset
        self.synaptic_duration = duration
    
    def set_channel(self,ID):
        prefix = self.ID_has_prefix(ID)
        type   = self.ID_is_type(ID)

        # first, make sure we haven't already set up this channel
        if ID in self.set_channels:
            raise RuntimeError('Already set up the channel with ID ' + str(ID))
        #-------- type 1 or type 5 (general type ion channel) --------------------
        if type == 1 or type == 5:
            # set up the right side of the m differential equation
            m_inf = lambda V: vd(self.channels_dict[prefix+'_m_inf'],V)
            tau_m = lambda V: vd(self.channels_dict[prefix+'_tau_m'],V)
            m_funct = lambda m,V: (m_inf(V) - m)/tau_m(V)

            # set up the right side of the h differential equation
            h_inf = lambda V: vd(self.channels_dict[prefix+'_h_inf'],V)
            tau_h = lambda V: vd(self.channels_dict[prefix+'_tau_h'],V)
            h_funct = lambda h,V: (h_inf(V) - h)/tau_h(V)

            self.set_channels[ID] = {'m_funct':m_funct,'h_funct':h_funct,'m_inf':m_inf,\
                                            'tau_m':tau_m,'h_inf':h_inf,'tau_h':tau_h} 
        #--------- type 2 (Sodium Current) ----------------------------------
        elif type == 2:
            # set up the right side of the m differential equation
            m_inf = lambda V: vd(self.channels_dict[prefix+'_m_inf'],V)
            tau_m = lambda V: vd(self.channels_dict[prefix+'_tau_m'],V)
            m_funct = lambda m,V: (m_inf(V) - m)/tau_m(V)

            # set up the right side of the h differential equation
            h_inf = lambda V: vd(self.channels_dict[prefix+'_h_inf'],V)
            # we have to change the way tau_h is calculated for this current
            tau_h1 = lambda V: vd(self.channels_dict[prefix+'_tau_h1'],V)
            tau_h2 = lambda V: vd(self.channels_dict[prefix+'_tau_h2'],V)
            tau_h = lambda V: tau_h1(V) * tau_h2(V)
            h_funct = lambda h,V: (h_inf(V) - h)/tau_h(V)

            self.set_channels[ID] = {'m_funct':m_funct,'h_funct':h_funct,'m_inf':m_inf,\
                                            'tau_m':tau_m,'h_inf':h_inf,'tau_h':tau_h} 
        #--------- type 3 (no h variable to calculate) ------------------------
        elif type == 3:
            # set up the right side of the m differential equation
            m_inf = lambda V: vd(self.channels_dict[prefix+'_m_inf'],V)
            tau_m = lambda V: vd(self.channels_dict[prefix+'_tau_m'],V)
            m_funct = lambda m,V: (m_inf(V) - m)/tau_m(V)

            # no h variable for this current

            self.set_channels[ID] = {'m_funct':m_funct,'m_inf':m_inf,\
                                            'tau_m':tau_m} 
        #--------- type 4 (Calcium Mediated Potassium Current) ------------
        elif type == 4:
            # set up the right side of the m differential equation
            m_inf = lambda Ca,V: (Ca/(Ca+3))*vd(self.channels_dict[prefix+'_m_inf'],V)
            tau_m = lambda V: vd(self.channels_dict[prefix+'_tau_m'],V)
            m_funct = lambda Ca,m,V: (m_inf(Ca,V) - m)/tau_m(V)

            # no h variable for this current

            self.set_channels[ID] = {'m_funct':m_funct,'m_inf':m_inf,\
                                            'tau_m':tau_m} 

def cell_data_setup(config_filename,sim_time):
    ''' Ensure that data for a given config_filename exists for the duration sim_time. '''

    # just to keep track of the configuration filename and the archiving directory name
    config = [ config_filename,config_filename[:-12] ] # _cell_config = 12 characters long

    # find out if the cell has a previous simulation run
    found_previous = False
    listing = os.listdir('.')
    for file in listing:
        if config[1] == file:
            c = Ccell.from_data_file( file+'/cell_data.cPickle')
            found_previous = True
    # ------------- run a novel simulation -----------------------
    if not found_previous:
        # set up the cell
        c = Ccell(config[0])
        
        # run the simulation for some time
        c.exercise(i_time = sim_time)

        # archive the resulting cell
        c.archive(config[1])
    else:
        #-------------- run more simulation time -----------------------
        # create the cell object from the data file
        c = Ccell.from_data_file(config[1] + '/cell_data.cPickle')

        if sim_time > c.total_time:
            # run the simulation for some ADDITIONAL time
            c.exercise(i_time = sim_time-c.total_time,X_init = c.X[-1])

            # archive the resulting cell
            c.archive(config[1])
        else:
            pass # we already have what data we requested.



    



    
 
